import os
import logging
import yaml
import datetime
from tqdm import tqdm

LOGTIME = datetime.datetime.now().strftime('%Y%m%d-%H-%M')


def exe_cmd(cmd):
    try:
        os.system(cmd)
    except Exception as e:
        logging.error(f"execute {cmd} wrong due to {e}")


def _load_yml(fpath, **kwargs_open):
    with open(fpath, **kwargs_open) as f:
        data = yaml.load(f.read(), Loader=yaml.FullLoader)
    return data


def load_yml(*args, **kwargs):
    if 'encoding' not in kwargs:
        for en in ['utf-8', None, 'gbk']:
            try:
                kwargs.update({'encoding': en})
                return _load_yml(*args, **kwargs)
            except:
                pass
    return _load_yml(*args, **kwargs)


def find_all_paths(target_path):
    paths = []
    for dirpath, dirname, filenames in os.walk(target_path):
        for filename in filenames:
            paths.append(os.path.join(dirpath, filename))
    return sorted(paths)


def get_all_ut_pys():
    all_pys = find_all_paths('tests/unit_tests/')
    ut_pys = list(filter(lambda x:  os.path.basename(x).startswith('test_') and x.endswith('.py'), all_pys))
    return ut_pys


def prepare_all(cmds):
    for script, cmd_dict in cmds['scripts'].items():
        for key, value in cmd_dict.items():
            if key == 'prepare':
                if isinstance(value, list):
                    for cmd in value:
                        exe_cmd(cmd)
                else:
                    for cmd in value.split('===cmdend==='):
                        exe_cmd(cmd.strip())


def create_pytest_script(path):
    run_script = f"../logs/pytest_scripts/{path.replace('tests/unit_tests/', '').replace(os.sep, '_')}"
    if not os.path.exists(os.path.dirname(run_script)):
        os.makedirs(os.path.dirname(run_script))
    template = """import os
import logging


def exe_cmd(cmd):
    try:
        os.system(cmd)
    except Exception as e:
        logging.error(f"execute {{cmd}} wrong due to {{e}}")


exe_cmd('pytest -s {}')
"""
    script_content = template.format(path)
    with open(run_script, 'w') as f:
        f.write(script_content)
    return run_script


if __name__ == '__main__':
    cmds = load_yml('commands.yml')
    ut_pys = get_all_ut_pys()
    # ut_pys = ['tests/unit_tests/models/test_gpt_model.py']
    # ut_pys = ut_pys[91:]

    prepare_all(cmds)

    with tqdm(total=len(ut_pys)) as _tqdm:
        for ut in tqdm(ut_pys):
            _tqdm.set_description(f'{ut} -')

            if ut in cmds['ignores']:
                print(f"ignore: {ut}, reason: {cmds['ignores'][ut]}")
                _tqdm.update(1)
                continue

            run_script = create_pytest_script(ut)

            nproc_per_node = 8
            if ut in cmds['scripts']:
                nproc_per_node = cmds['scripts'][ut].get('nproc_per_node', nproc_per_node)
            
            log_dir = f"../logs/utlogs/{LOGTIME}"
            if not os.path.exists(log_dir):
                os.makedirs(log_dir)
            log_path = f"{log_dir}/{ut.replace('tests/unit_tests/', '').replace(os.sep, '_').replace('.py', '.log')}"
            runcmd = "python -m torch.distributed.launch --master_addr=localhost --master_port=6000 --nproc_per_node={} --nnodes=1 --node_rank=0 {} > {} 2>&1".format(nproc_per_node, run_script, log_path)
            exe_cmd(runcmd)
            _tqdm.update(1)
